// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#pragma once

#include <melon/socket_address.h>
#include <melon/io/async/event_base_thread.h>
#include <pollux/plan/reference_query_runner.h>
#include <pollux/vector/complex_vector.h>

using namespace kumo::pollux::plan;

namespace kumo::pollux::exec::test {
    template<typename T>
    T extractSingleValue(const std::vector<RowVectorPtr> &data) {
        auto simpleVector = data[0]->childAt(0)->as<SimpleVector<T> >();
        POLLUX_CHECK(!simpleVector->is_null_at(0));
        return simpleVector->value_at(0);
    }

    class QueryRunnerContext {
    public:
        melon::F14FastMap<core::PlanNodeId, std::vector<std::string> > windowFrames_;
    };

    class PrestoQueryRunner : public ReferenceQueryRunner {
    public:
        /// @param coordinatorUri Presto REST API endpoint, e.g. http://127.0.0.1:8080
        /// @param user Username to use in X-Presto-User header.
        /// @param timeout Timeout in milliseconds of an HTTP request.
        PrestoQueryRunner(
            memory::MemoryPool *aggregatePool,
            std::string coordinatorUri,
            std::string user,
            std::chrono::milliseconds timeout);

        RunnerType runnerType() const override {
            return RunnerType::kPrestoQueryRunner;
        }

        const std::vector<TypePtr> &supportedScalarTypes() const override;

        const std::unordered_map<std::string, DataSpec> &
        aggregationFunctionDataSpecs() const override;

        /// Converts Pollux query plan to Presto SQL. Supports Values -> Aggregation or
        /// Window with an optional Project on top.
        ///
        /// Values node is converted into reading from 'tmp' table.
        ///
        /// @return std::nullopt if Values node uses types not supported by DWRF file
        /// format (DATE, INTERVAL, UNKNOWN).
        std::optional<std::string> toSql(
            const pollux::core::PlanNodePtr &plan) override;

        bool isConstantExprSupported(const core::TypedExprPtr &expr) override;

        bool isSupported(const exec::FunctionSignature &signature) override;

        // Converts 'plan' into an SQL query and executes it. Result is returned as a
        // MaterializedRowMultiset with the ReferenceQueryErrorCode::kSuccess if
        // successful, or an std::nullopt with a ReferenceQueryErrorCode if the query
        // fails.
        std::pair<
            std::optional<std::multiset<std::vector<pollux::variant> > >,
            ReferenceQueryErrorCode>
        execute(const core::PlanNodePtr &plan) override;

        /// Similar to 'execute' but returns results in RowVector format.
        /// Caller should ensure 'supportsPolluxVectorResults' returns true.
        std::pair<
            std::optional<std::vector<pollux::RowVectorPtr> >,
            ReferenceQueryErrorCode>
        executeAndReturnVector(const core::PlanNodePtr &plan) override;

        /// Executes Presto SQL query and returns the results. Tables referenced by
        /// the query must already exist.
        std::vector<pollux::RowVectorPtr> execute(const std::string &sql) override;

        /// Executes Presto SQL query with extra presto session property.
        std::vector<pollux::RowVectorPtr> execute(
            const std::string &sql,
            const std::string &sessionProperty) override;

        bool supportsPolluxVectorResults() const override;

        std::shared_ptr<QueryRunnerContext> queryRunnerContext() {
            return queryRunnerContext_;
        }

    private:
        using ReferenceQueryRunner::toSql;

        memory::MemoryPool *pool() {
            return pool_.get();
        }

        std::string startQuery(
            const std::string &sql,
            const std::string &sessionProperty = "");

        std::string fetchNext(const std::string &nextUri);

        // Creates an empty table with given data type and table name. The function
        // returns the root directory of table files.
        std::string createTable(const std::string &name, const TypePtr &type);

        const std::string coordinatorUri_;
        const std::string user_;
        const std::chrono::milliseconds timeout_;
        melon::EventBaseThread eventBaseThread_{false};
        std::shared_ptr<memory::MemoryPool> pool_;
        std::shared_ptr<QueryRunnerContext> queryRunnerContext_;
    };
} // namespace kumo::pollux::exec::test
